import torch
import torchvision
from torch import nn
from torch.nn import Conv2d
from torch.utils.data import DataLoader

datasets = torchvision.datasets.CIFAR10(root='./data/data_non_linear', train=False, download=True,
                                        transform=torchvision.transforms.ToTensor())

train_loader = DataLoader(datasets, batch_size=1)


class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()

        self.model1 = nn.Sequential(
            Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(kernel_size=2),
            Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(kernel_size=2),
            Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = self.model1(x)
        return x


loss = nn.CrossEntropyLoss()
net = SimpleNN()
optim = torch.optim.SGD(net.parameters(), lr=0.01)
for epoch in range(20):
    running_loss = 0.0
    for data in train_loader:
        inputs, labels = data
        outputs = net(inputs)
        # print(f"输出output：{outputs}")
        # print(f"标签labels：{labels}")
        result_loss = loss(outputs, labels)
        optim.zero_grad()
        result_loss.backward()
        optim.step()
        # print(f"第{epoch + 1}轮——loss:{result_loss}")
        running_loss += result_loss

    print(f"第{epoch + 1}轮——loss:{running_loss}")
